# %%
# import packages
from __future__ import print_function, division

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
from torch.utils.data import Dataset, DataLoader
import torchvision
from torchvision import datasets, models, transforms
import matplotlib.pyplot as plt
import time
import os
import copy
import cv2

# plt.ion()   # interactive mode
# %%
# dataset
data_dir = "D:\\CS\\SRC\\pyFoo\\pytorch\\data\\us_image\\origin"

# data_transforms = transforms.Compose([
#     transforms.Resize(256),
#     # transforms.RandomResizedCrop(224),
#     transforms.RandomHorizontalFlip(),
#     transforms.RandomVerticalFlip(),
#     transforms.RandomRotation(True),
#     transforms.ToTensor(),
#     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])



data_transforms = transforms.Compose([
    # transforms.Resize(256),
    # # transforms.RandomResizedCrop(224),
    # transforms.RandomHorizontalFlip(),
    # transforms.RandomVerticalFlip(),
    # transforms.RandomRotation(True),
    transforms.ToTensor(),
    # transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])    
    ])



# {
    # 'trian': transforms.Compose([
    #     transforms.Resize(256),
    #     transforms.RandomResizedCrop(224),
    #     transforms.RandomHorizontalFlip(),
    #     transforms.RandomVerticalFlip(),
    #     transforms.RandomRotation(True),
    #     transforms.ToTensor(),
    #     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    # ]),
    # 'val': transforms.Compose([
    #     transforms.Resize(256),
    #     transforms.CenterCrop(224),
    #     transforms.ToTensor(),
    #     transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    # ]),
# }


# image_datasets = {x: datasets.ImageFolder(os.path.join(data_dir, x),
#                                           data_transforms[x])
#                   for x in ['trian', 'val']}


data_dir="D:\\CS\\SRC\\pyFoo\\pytorch\\data\\us_image\\train"
image_datasets=datasets.ImageFolder(
    data_dir, data_transforms)

# class_names = image_datasets['train'].classes

device=torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print("using device of ", device)

# %%
# verify loading image

print("image_datasets length is ", image_datasets.__len__())
image_datasets.class_to_idx
# image_datasets.imgs
im_tensor, label = image_datasets[0]

# 转化为numpy，需要opencv显示
im_num = im_tensor.numpy()*255
im_num = im_num.astype('uint8')
im_num = np.transpose(im_num, (1, 2, 0))

# 转化为PILImage
im_pil = transforms.ToPILImage()(im_tensor).convert('L')
# im_pil = transforms.ToPILImage()(im_tensor)
print("pil image size is ", im_pil.size)
print("pil image mode is ", im_pil.mode)


img_origin = cv2.imread("D:\\CS\\SRC\\pyFoo\\pytorch\\data\\us_image\\train\\origin\\1.1.jpg")
# cv2.namedWindow("Image")
# cv2.imshow("Image", img_origin)
# cv2.waitKey(0)
# cv2.destroyAllWindows()

# PILImage Show
plt.figure()
plt.subplot(311)
# plt.axis('off')
plt.title('im_pil')
plt.imshow(im_pil, cmap='gray')

plt.subplot(312)
# plt.axis('off')
plt.title('im_num')
plt.imshow(im_num, cmap='gray')

plt.subplot(313)
plt.title('img_origin')
plt.imshow(img_origin, cmap='gray')
# im_origin = os.path.join(data_dir, "1.1.jpg")


# %%
# dataload

dataset_loader = DataLoader(image_datasets, batch_size=4, shuffle=True, num_workers=4)


#%%
# Helper function to show a batch
def show_batch(sample_batched):
    """Show image with landmarks for a batch of samples."""
    images_batch, landmarks_batch = \
        sample_batched['image'], sample_batched['landmarks']
    batch_size = len(images_batch)
    im_size = images_batch.size(2)

    grid = utils.make_grid(images_batch)
    plt.imshow(grid.numpy().transpose((1, 2, 0)))


for i_batch, sample_batched in enumerate(dataset_loader):
    print(i_batch, sample_batched['image'].size(),
          sample_batched['landmarks'].size())

    # observe 4th batch and stop.
    if i_batch == 3:
        plt.figure()
        show_landmarks_batch(sample_batched)
        plt.axis('off')
        plt.ioff()
        plt.show()
        break
